import numpy as np
import json


def get_index_cmd(ds_name, build_graph_para, project_dir, function_type):
    if function_type == 'l2':
        pre_cmd = './tests/test_nndescent_l2'
    elif function_type == 'inner_product':
        pre_cmd = './tests/test_nndescent_inner_product'
    elif function_type == 'cosine':
        pre_cmd = './tests/test_nndescent_cosine'
    else:
        raise Exception('not support function type')

    base_fname = '%s/data/%s/%s_base.fvecs' % (project_dir, ds_name, ds_name)
    graph_fname = '%s/graph/%s-%dNN-%s.knng' % (project_dir, ds_name, build_graph_para['K'], function_type)
    command = '%s %s %s %d %d %d %d %d' % (
        pre_cmd, base_fname, graph_fname, build_graph_para['K'], build_graph_para['L'], build_graph_para['iter'],
        build_graph_para['S'], build_graph_para['R'])
    return command


if __name__ == "__main__":
    para_K = 100
    efsearch = 200
    n_iter = 10
    para_S = 15
    para_R = 100
    build_graph_para = {
        'K': para_K,
        'L': efsearch,
        'iter': n_iter,
        'S': para_S,
        'R': para_R
    }
    project_dir = '/home/zhengbian/knn-graph'
    ds_name_l = ['audio', 'imagenet', 'movielens', 'music100', 'netflix', 'normal-64', 'text-to-image', 'tiny5m',
                 'word2vec', 'yahoomusic']
    cmd_m = {}
    for dataset_name in ds_name_l:
        l2_cmd = get_index_cmd(dataset_name, build_graph_para, project_dir, 'l2')
        cosine_cmd = get_index_cmd(dataset_name, build_graph_para, project_dir, 'cosine')
        ip_cmd = get_index_cmd(dataset_name, build_graph_para, project_dir, 'inner_product')
        cmd_m[dataset_name] = {
            'l2': l2_cmd,
            'cosine': cosine_cmd,
            'inner_product': ip_cmd
        }
    with open('command.json', 'w') as f:
        json.dump(cmd_m, f)
